package com.demo.wrapper;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.map.MapUtil;
import co.elastic.clients.elasticsearch._types.SortOptions;
import co.elastic.clients.elasticsearch._types.SortOrder;
import co.elastic.clients.elasticsearch._types.aggregations.Aggregation;
import co.elastic.clients.elasticsearch._types.query_dsl.BoolQuery;
import co.elastic.clients.elasticsearch._types.query_dsl.ChildScoreMode;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.search.SourceFilter;
import co.elastic.clients.json.JsonData;
import co.elastic.clients.util.ObjectBuilder;
import com.demo.wrapper.enums.BoolEnum;
import com.demo.wrapper.enums.OperatorEnum;
import com.demo.wrapper.enums.OrderEnum;
import com.demo.wrapper.enums.ScoreModeEnum;
import com.demo.wrapper.pojo.Agg;
import com.demo.wrapper.pojo.Arg;
import com.demo.wrapper.pojo.ConditionalReceiver;
import com.demo.wrapper.pojo.Order;
import com.demo.wrapper.utils.ArrayUtil;
import com.demo.wrapper.utils.FieldValueUtil;

import java.util.*;
import java.util.function.Function;

/**
 * es条件拼接器处理
 *
 * @author 19242
 */
public class WrapperUtil {
    /**
     * 将wrapper收集到的条件转换成QueryBuilder
     *
     * @param wrapper
     * @param autoFilter
     * @param <T>
     * @return
     */
    public static <T extends AbstractEsQueryWrapper<T, ?, Query.Builder, SearchRequest.Builder>> Query.Builder getBuilder(T wrapper, boolean autoFilter) {
        if (CollectionUtil.isEmpty(wrapper.getChildren(false))) {
            Set<Arg> querySet = wrapper.getQuerySet(false);
            if (CollectionUtil.isEmpty(querySet)) {
                return null;
            }
            // 如果没有子查询且只有一条查询条件，是and或or关系的，返回单查询
            if (autoFilter && 1 == querySet.size()) {
                if (BoolEnum.MUST == wrapper.getBoolEnum() || BoolEnum.SHOULD == wrapper.getBoolEnum()) {
                    return WrapperUtil.setQueryBuilder(querySet.stream().findFirst().orElse(null), true);
                }
            }
        }
        return getBuilder(wrapper, autoFilter, null);
    }

    /**
     * 将wrapper收集到的条件转换成QueryBuilder
     *
     * @param wrapper       解析的拼接器
     * @param autoFilter    自动过滤相同关系的子条件，不支持not
     * @param parentBuilder 父节点的QueryBuilder
     * @param <T>           泛型
     * @return QueryBuilder
     */
    static <T extends AbstractEsQueryWrapper<T, ?, Query.Builder, SearchRequest.Builder>> Query.Builder getBuilder(T wrapper, boolean autoFilter, BoolQuery.Builder parentBuilder) {
        BoolQuery.Builder boolQueryBuilder = new BoolQuery.Builder();
        Map<BoolEnum, List<T>> children = wrapper.getChildren(false);
        if (MapUtil.isNotEmpty(children)) {
            // not的不支持过滤
            if (autoFilter && wrapper.getBoolEnum() != BoolEnum.MUST_NOT) {
                // 当前节点和子节点是相同关系，把子节点条件并入当前条件
                List<T> common = children.remove(wrapper.getBoolEnum());
                mergeCommonChildren(wrapper, common);
            }
            // 不同关系的子节点收集起来
            for (Map.Entry<BoolEnum, List<T>> entry : children.entrySet()) {
                for (T esQueryWrapper : entry.getValue()) {
                    setBoolBuilder(boolQueryBuilder, esQueryWrapper.getBoolEnum(), getBuilder(esQueryWrapper, autoFilter, boolQueryBuilder));
                }
            }
        }
        if (CollectionUtil.isNotEmpty(wrapper.getQuerySet(false))) {
            // 优先使用父节点拼接条件
            BoolQuery.Builder builder = Objects.nonNull(parentBuilder) ? parentBuilder : boolQueryBuilder;
            for (Arg arg : wrapper.getQuerySet(false)) {
                Query.Builder query = setQueryBuilder(arg, autoFilter);
                if (Objects.isNull(query)) {
                    continue;
                }
                setBoolBuilder(builder, wrapper.getBoolEnum(), query);
            }
        }
        return nonEmpty(boolQueryBuilder);
    }

    /**
     * 条件为空 返回null
     *
     * @param boolQueryBuilder
     * @return
     */
    private static Query.Builder nonEmpty(BoolQuery.Builder boolQueryBuilder) {
        if (Objects.isNull(boolQueryBuilder)) {
            return null;
        }
        BoolQuery boolQuery = boolQueryBuilder.build();
        if (CollectionUtil.isEmpty(boolQuery.must()) && CollectionUtil.isEmpty(boolQuery.should()) && CollectionUtil.isEmpty(boolQuery.mustNot()) && CollectionUtil.isEmpty(boolQuery.filter())) {
            return null;
        }
        Query.Builder queryBuilder = new Query.Builder();
        queryBuilder.bool(boolQuery);
        return queryBuilder;
    }

    /**
     * 合并相同关系的子查询
     *
     * @param wrapper 当前节点
     * @param common  关系相同的子节点
     */
    private static <Wrapper extends AbstractEsQueryWrapper<Wrapper, ?, Query.Builder, SearchRequest.Builder>> void mergeCommonChildren(Wrapper wrapper, List<Wrapper> common) {
        if (CollectionUtil.isEmpty(common)) {
            return;
        }
        for (Wrapper esQueryWrapper : common) {
            // 将子节点的条件放入当前节点
            wrapper.addAllQuery(esQueryWrapper.getQuerySet(false));
            Map<BoolEnum, List<Wrapper>> children = esQueryWrapper.getChildren(false);
            if (CollectionUtil.isNotEmpty(children)) {
                List<Wrapper> remove = children.remove(wrapper.getBoolEnum());
                // 遍历子节点的子节点添加相同节点
                mergeCommonChildren(wrapper, remove);
                // 清空子节点的查询条件
                esQueryWrapper.clearQuery();
                // 将子查询放回
                wrapper.addChildren(esQueryWrapper);
            }
        }
    }

    /**
     * 添加关系
     *
     * @param boolQueryBuilder 多条件拼接
     * @param boolEnum         关系类型
     * @param queryBuilder     条件
     */
    static void setBoolBuilder(BoolQuery.Builder boolQueryBuilder, BoolEnum boolEnum, Query.Builder queryBuilder) {
        if (Objects.isNull(queryBuilder)) {
            return;
        }
        Query query = queryBuilder.build();
        switch (boolEnum) {
            case MUST:
                boolQueryBuilder.must(query);
                break;
            case SHOULD:
                boolQueryBuilder.should(query);
                break;
            case MUST_NOT:
                boolQueryBuilder.mustNot(query);
                break;
            case FILTER:
                boolQueryBuilder.filter(query);
            default:
                break;
        }
    }

    /**
     * 拼接查询条件
     *
     * @param arg        查询参数
     * @param autoFilter
     * @return 条件
     */
    static Query.Builder setQueryBuilder(Arg arg, boolean autoFilter) {
        if (Objects.isNull(arg) || Objects.isNull(arg.getQueryEnum())) {
            return null;
        }
        Query.Builder queryBuilder = new Query.Builder();
        switch (arg.getQueryEnum()) {
            case TERM:
                queryBuilder.term(v -> v.field(arg.getColumn()).value(FieldValueUtil.fieldValue(arg.getVal())).boost(arg.getBoost()));
                break;
            case TERMS:
                queryBuilder.terms(v -> v.field(arg.getColumn()).terms(FieldValueUtil.termsQueryField(arg.getVal())).boost(arg.getBoost()));
                break;
            case MATCH:
                queryBuilder.match(v -> v.field(arg.getColumn()).query(FieldValueUtil.fieldValue(arg.getVal())).boost(arg.getBoost()).operator(FieldValueUtil.operator((OperatorEnum) arg.getExt1())));
                break;
            case NESTED:
                queryBuilder.nested(v -> v.path((String) arg.getExt1()).query(((AbstractEsQueryWrapper<?, ?, Query.Builder, ?>) arg.getExt2()).queryBuilder(autoFilter).build()).scoreMode(getScoreMode((ScoreModeEnum) arg.getExt3())));
                break;
            case MATCH_ALL:
                queryBuilder.matchAll(v -> v.boost(arg.getBoost()));
                break;
            case MATCH_PHRASE:
                queryBuilder.matchPhrase(v -> v.field(arg.getColumn()).query((String) arg.getVal()).boost(arg.getBoost()));
                break;
            case MATCH_PHRASE_PREFIX:
                queryBuilder.matchPhrasePrefix(v -> v.field(arg.getColumn()).query((String) arg.getVal()).boost(arg.getBoost()).maxExpansions((int) arg.getExt1()));
                break;
            case MULTI_MATCH:
                queryBuilder.multiMatch(v -> v.query((String) arg.getVal()).fields(Arrays.asList((String[]) arg.getExt1())));
                break;
            case GT:
                queryBuilder.range(v -> v.field(arg.getColumn()).gt(JsonData.of(arg.getVal())).boost(arg.getBoost()));
                break;
            case GE:
                queryBuilder.range(v -> v.field(arg.getColumn()).gte(JsonData.of(arg.getVal())).boost(arg.getBoost()));
                break;
            case LT:
                queryBuilder.range(v -> v.field(arg.getColumn()).lt(JsonData.of(arg.getVal())).boost(arg.getBoost()));
                break;
            case LE:
                queryBuilder.range(v -> v.field(arg.getColumn()).lte(JsonData.of(arg.getVal())).boost(arg.getBoost()));
                break;
            case BETWEEN:
                queryBuilder.range(v -> v.field(arg.getColumn()).gte(JsonData.of(arg.getExt1())).lte(JsonData.of(arg.getExt2())).boost(arg.getBoost()));
                break;
            case WILDCARD:
                queryBuilder.wildcard(v -> v.field(arg.getColumn()).value((String) arg.getVal()).boost(arg.getBoost()));
                break;
            case EXISTS:
                queryBuilder.exists(v -> v.field(arg.getColumn()).boost(arg.getBoost()));
                break;
            case QUERY_BUILDER:
                ((Function<Query.Builder, ObjectBuilder<Query>>) arg.getExt1()).apply(queryBuilder);
                break;
            default:
                break;
        }
        return queryBuilder;
    }

    /**
     * 获取嵌套聚合得分模式
     *
     * @param scoreModeEnum
     * @return
     */
    private static ChildScoreMode getScoreMode(ScoreModeEnum scoreModeEnum) {
        if (Objects.isNull(scoreModeEnum)) {
            return ChildScoreMode.Avg;
        }
        switch (scoreModeEnum) {
            case Avg:
                return ChildScoreMode.Avg;
            case Max:
                return ChildScoreMode.Max;
            case Min:
                return ChildScoreMode.Min;
            case Sum:
                return ChildScoreMode.Sum;
            default:
                return ChildScoreMode.None;
        }
    }

    /**
     * 获取SearchSourceBuilder
     *
     * @param <Wrapper>  AbstractEsQueryWrapper的子类
     * @param wrapper    条件
     * @param autoFilter
     * @return
     */
    static <Wrapper extends AbstractEsQueryWrapper<Wrapper, ?, Query.Builder, SearchRequest.Builder>> SearchRequest.Builder getSourceBuilder(Wrapper wrapper, boolean autoFilter) {
        SearchRequest.Builder searchSourceBuilder = new SearchRequest.Builder();
        Query.Builder queryBuilder = wrapper.queryBuilder(autoFilter);
        if (Objects.nonNull(queryBuilder)) {
            searchSourceBuilder.query(queryBuilder.build());
        }
        List<Arg> limitList = wrapper.getLimitList(false);
        if (CollectionUtil.isNotEmpty(limitList)) {
            // 防止有些条件被覆盖掉，先收集完再处理
            ConditionalReceiver receiver = new ConditionalReceiver();
            for (Arg arg : limitList) {
                setLimit(searchSourceBuilder, arg, receiver);
            }
            if (CollectionUtil.isNotEmpty(receiver.getIncludes()) || CollectionUtil.isNotEmpty(receiver.getExcludes())) {
                SourceFilter.Builder builder = new SourceFilter.Builder();
                if (CollectionUtil.isNotEmpty(receiver.getIncludes())) {
                    builder.includes(receiver.getIncludes());
                }
                if (CollectionUtil.isNotEmpty(receiver.getExcludes())) {
                    builder.excludes(receiver.getExcludes());
                }
                searchSourceBuilder.source(v -> v.filter(builder.build()));
            }
        }
        return searchSourceBuilder;
    }

    /**
     * 添加条件
     *
     * @param searchSourceBuilder
     * @param arg
     * @return
     */
    private static void setLimit(SearchRequest.Builder searchSourceBuilder, Arg arg, ConditionalReceiver receiver) {
        if (Objects.isNull(arg) || Objects.isNull(arg.getQueryEnum())) {
            return;
        }
        switch (arg.getQueryEnum()) {
            case FROM:
                searchSourceBuilder.from((int) arg.getExt1());
                break;
            case SIZE:
                searchSourceBuilder.size((int) arg.getExt1());
                break;
            case LIMIT:
                searchSourceBuilder.from((int) arg.getExt1()).size((int) arg.getExt2());
                break;
            case INCLUDES:
                ArrayUtil.addAll(receiver.getIncludes(), (String) arg.getExt1(), (String[]) arg.getExt2());
                break;
            case EXCLUDES:
                ArrayUtil.addAll(receiver.getExcludes(), (String) arg.getExt1(), (String[]) arg.getExt2());
                break;
            case FETCH_SOURCE:
                ArrayUtil.addAll(receiver.getIncludes(), (String[]) arg.getExt1());
                ArrayUtil.addAll(receiver.getExcludes(), (String[]) arg.getExt2());
                break;
            case AGG:
                searchSourceBuilder.aggregations((String) arg.getExt1(), (Aggregation) arg.getExt2());
                break;
            case SORT:
                searchSourceBuilder.sort(SortOptions.of(v -> v.field(c -> c.field(arg.getColumn()).order(FieldValueUtil.sortOrder((OrderEnum) arg.getExt1())))));
                break;
            case SOURCE_BUILDER:
                ((Function<SearchRequest.Builder, SearchRequest.Builder>) arg.getExt1()).apply(searchSourceBuilder);
                break;
            default:
                break;
        }
    }


    /**
     * 组装聚合条件
     *
     * @param builder
     * @param agg
     * @return
     */
    public static ObjectBuilder<Aggregation> setAgg(Aggregation.Builder builder, Agg agg) {
        if (Objects.isNull(agg) || Objects.isNull(agg.getAggEnum())) {
            return null;
        }
        switch (agg.getAggEnum()) {
            case TERMS:
                Order[] orders = (Order[]) agg.getExt2();
                List<Map<String, SortOrder>> sortOrderList = new ArrayList<>();
                for (Order order : orders) {
                    sortOrderList.add(MapUtil.of(order.getColumn(), FieldValueUtil.sortOrder(order.isAsc())));
                }
                return builder.terms(v -> v.field(agg.getColumn()).size((int) agg.getExt1()).order(sortOrderList));
            case VALUE_COUNT:
                return builder.valueCount(v -> v.field(agg.getColumn()));
            case CARDINALITY:
                return builder.cardinality(v -> v.field(agg.getColumn()));
            case AVG:
                return builder.avg(v -> v.field(agg.getColumn()));
            case MAX:
                return builder.max(v -> v.field(agg.getColumn()));
            case MIN:
                return builder.min(v -> v.field(agg.getColumn()));
            case SUM:
                return builder.sum(v -> v.field(agg.getColumn()));
            case STATS:
                return builder.stats(v -> v.field(agg.getColumn()));
            case BUILDER:
                return ((Function<Aggregation.Builder, ObjectBuilder<Aggregation>>) agg.getExt1()).apply(builder);
            case FILTER:
                Query.Builder queryBuilder = ((AbstractEsQueryWrapper<?, ?, Query.Builder, ?>) agg.getExt1()).queryBuilder();
                return builder.filter(queryBuilder.build());
            default:
                return null;
        }
    }

    /**
     * 将收集的Agg转成AggregationBuilder
     *
     * @return AggregationBuilder
     */
    public static Aggregation getAggregationBuilder(AbstractEsAggWrapper<? extends AbstractEsAggWrapper<?, ?>, ?> aggWrapper) {
        return Aggregation.of(v -> {
            List<? extends AbstractEsAggWrapper<? extends AbstractEsAggWrapper<?, ?>, ?>> aggChildren = aggWrapper.getAggChildren(false);
            if (CollectionUtil.isNotEmpty(aggChildren)) {
                for (AbstractEsAggWrapper<? extends AbstractEsAggWrapper<?, ?>, ?> aggChild : aggChildren) {
                    Optional.ofNullable(getAggregationBuilder(aggChild)).ifPresent(b -> v.aggregations(aggChild.getAgg().getName(), b));
                }
            }
            return WrapperUtil.setAgg(v, aggWrapper.getAgg());
        });
    }

}
